use std::os::raw::c_void;

use melior::{
    ir::{OperationRef, TypeLike, ValueLike, operation::OperationLike, r#type::TypeId},
    pass::{ExternalPass, Pass, create_external},
};
use mlir_sys::{
    MlirOperation, MlirWalkOrder_MlirWalkPostOrder, MlirWalkResult_MlirWalkResultAdvance,
    mlirOperationWalk,
};

use crate::rust_side::mlir_dyn_cast_shape_inference;

#[repr(align(8))]
struct ShapeInferencePassId;

static SHAPE_INFERENCE_PASS: ShapeInferencePassId = ShapeInferencePassId;

pub fn create_shape_inference() -> Pass {
    create_external(
        shape_inference_pass,
        TypeId::create(&SHAPE_INFERENCE_PASS),
        "ShapeInference",
        "toy-shape-inference",
        "Interface to access a registered method to infer the return types for an operation that can be used during type inference.",
        "toy.func",
        &[],
    )
}

fn shape_inference_pass(op_ref: OperationRef, pass: ExternalPass) {
    let mut op_work_list = Vec::<MlirOperation>::new();
    unsafe extern "C" fn callback(raw_op: MlirOperation, data: *mut c_void) -> u32 {
        if returns_dynamic_shape(raw_op) {
            let data = data as *mut Vec<MlirOperation>;
            let list: &mut Vec<MlirOperation> = unsafe { &mut *data };
            list.push(raw_op);
        }
        return MlirWalkResult_MlirWalkResultAdvance;
    }
    let raw_op = op_ref.to_raw();
    unsafe {
        mlirOperationWalk(
            raw_op,
            Some(callback),
            &mut op_work_list as *mut _ as *mut c_void,
            MlirWalkOrder_MlirWalkPostOrder,
        );
    }

    while !op_work_list.is_empty() {
        let next_op_index = op_work_list
            .iter()
            .position(|&raw_op| all_operands_inferred(raw_op));

        if next_op_index.is_none() {
            break;
        }

        let raw_op = op_work_list.swap_remove(next_op_index.unwrap());
        let op_ref = unsafe { OperationRef::from_raw(raw_op) };

        if let Some(shape_op) = mlir_dyn_cast_shape_inference(op_ref) {
            shape_op.infer_shapes();
        } else {
            eprintln!("unable to infer shape");
            pass.signal_failure();
        }
    }

    if !op_work_list.is_empty() {
        eprintln!("shape infer failed");
        pass.signal_failure();
    }
}

fn all_operands_inferred(raw_op: MlirOperation) -> bool {
    let op_ref = unsafe { OperationRef::from_raw(raw_op) };
    op_ref
        .operands()
        .all(|operand| operand.r#type().is_ranked_tensor())
}

fn returns_dynamic_shape(raw_op: MlirOperation) -> bool {
    let op_ref = unsafe { OperationRef::from_raw(raw_op) };

    op_ref
        .results()
        .any(|result| !result.r#type().is_ranked_tensor())
}
